#!/usr/bin/python3
# -*- coding: utf-8; tab-width: 4; indent-tabs-mode: t -*-

import os
import sys
import shutil
import pyudev
import parted
import subprocess


class Util:

    @staticmethod
    def isMbrOrGpt(devPath):
        # Use parted to get disk information
        dev = parted.getDevice(devPath)
        disk = parted.newDisk(dev)
        if disk.type == "msdos":
            return True
        elif disk.type == "gpt":
            return False
        else:
            assert False

    @staticmethod
    def forceDelete(path):
        if os.path.islink(path):
            os.remove(path)
        elif os.path.isfile(path):
            os.remove(path)
        elif os.path.isdir(path):
            shutil.rmtree(path)
        elif os.path.lexists(path):
            os.remove(path)             # other type of file, such as device node
        else:
            pass                        # path does not exist, do nothing

    @staticmethod
    def die(message):
        print(message, file=sys.stderr)
        sys.exit(1)

    @staticmethod
    def getDiskUsbVendorModel(devPath):
        ctx = pyudev.Context()
        deviceObj = pyudev.Devices.from_device_file(ctx, devPath)
        return (deviceObj.properties["ID_USB_VENDOR_ID"], deviceObj.properties["ID_USB_MODEL_ID"])


class Handy:

    @staticmethod
    def usage():
        msg = ""
        msg += 'Usage: %s install <assist-disk-path> <--bios|--efi>\n' % (sys.argv[0])
        msg += '       %s run' % (sys.argv[0])
        return msg


if __name__ == "__main__":
    selfDir = os.path.dirname(os.path.realpath(__file__))
    imageFullfn = os.path.join(selfDir, "disk.img")

    if len(sys.argv) < 2:
        Util.die(Handy.usage())
    testType = sys.argv[1]

    if testType == "install":
        if len(sys.argv) < 4:
            Util.die(Handy.usage())
        testFullfn = sys.argv[2]
        if sys.argv[3] not in ["--bios", "--efi"]:
            Util.die(Handy.usage())

        subprocess.check_call(["dd", "if=/dev/zero", "of=%s" % (imageFullfn), "bs=%d" % (1024 * 1024), "count=%d" % (4096), "conv=sparse"], stderr=subprocess.DEVNULL)

        vmCmd = "qemu-system-x86_64"
        vmCmd += " -no-user-config -nodefaults"
        vmCmd += " -machine q35,usb=on,accel=kvm"
        vmCmd += " -cpu host -smp 2,sockets=1,cores=2,threads=1"
        vmCmd += " -m 4096M"
        vmCmd += " -rtc base=localtime"
        if sys.argv[3] == "efi":
            vmCmd += " -drive if=pflash,format=raw,readonly=on,file=/usr/share/edk2/OvmfX64/OVMF_CODE.fd"
        vmCmd += " -device isa-fdc"
        vmCmd += " -device VGA"
        vmCmd += " -device usb-tablet"
        # if True:
        #     vmCmd += " -blockdev 'node-name=main-disk,driver=file,filename=%s,cache.no-flush=on'" % (imageFullfn)
        #     vmCmd += " -device ide-hd,bus=ide.0,unit=0,drive=main-disk"
        if True:
            if testFullfn.startswith("/dev/"):
                vendorId, modelId = Util.getDiskUsbVendorModel(testFullfn)
                vmCmd += " -device usb-host,vendorid=0x%s,productid=0x%s" % (vendorId, modelId)
            else:
                vmCmd += " -blockdev 'node-name=usb-stick,driver=file,filename=%s,cache.no-flush=on'" % (testFullfn)
                vmCmd += " -device usb-storage,drive=usb-stick,removable=on"
        vmCmd += " -boot order=n,menu=on"
        subprocess.check_call(vmCmd, shell=True)

        sys.exit(0)

    if testType == "run":
        if not os.path.exists(imageFullfn):
            Util.die("Image file does not exist.")

        vmCmd = "qemu-system-x86_64"
        vmCmd += " -no-user-config -nodefaults"
        vmCmd += " -machine q35,usb=on,accel=kvm"
        vmCmd += " -cpu host -smp 2,sockets=1,cores=2,threads=1"
        vmCmd += " -m 4096M"
        vmCmd += " -rtc base=localtime"
        if not Util.isMbrOrGpt(imageFullfn):
            vmCmd += " -drive if=pflash,format=raw,readonly=on,file=/usr/share/edk2/OvmfX64/OVMF_CODE.fd"
        vmCmd += " -device isa-fdc"
        vmCmd += " -device VGA"
        vmCmd += " -device usb-tablet"
        if True:
            vmCmd += " -blockdev 'node-name=main-disk,driver=file,filename=%s,cache.no-flush=on'" % (imageFullfn)
            vmCmd += " -device ide-hd,bus=ide.0,unit=0,drive=main-disk"
        subprocess.check_call(vmCmd, shell=True)

        sys.exit(0)

    Util.die('Test type should be "install" or "run")')